import os.path
import numpy as np
import os
import os.path
import random
import h5py
import torchvision.transforms as transforms
import cv2
import glob
import torch.utils.data as udata


def normalize(data):
    return data/255.


def Im2Patch(img, win, stride):
    k = 0
    endw = img.shape[0]
    endh = img.shape[1]
    endc = img.shape[2]
    patch = img[0:endw-win+0+1:stride, 0:endh-win+0+1:stride, :]
    TotalPatNum = patch.shape[0] * patch.shape[1]
    Y = np.zeros([endc, win*win, TotalPatNum], np.uint8)
    for i in range(win):
        for j in range(win):
            patch = img[i:endw-win+i+1:stride, j:endh-win+j+1:stride, :]
            Y[:, k, :] = np.array(patch[:]).reshape(endc, TotalPatNum)
            k = k + 1
    Y = Y.reshape([endc, win, win, TotalPatNum])
    return np.transpose(Y, [1, 2, 0, 3])



def prepare_data(filename, data_path, patch_size, stride):
    # train
    print('process data')
    # scales = [1, 0.9, 0.8]
    scales = [1]
    files = glob.glob(os.path.join(data_path, '*.*'))
    files.sort()
    h5f = h5py.File(filename, 'w')
    train_num = 0
    for i in range(len(files)):
        img = cv2.imread(files[i])
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # h, w = img.shape
        for k in range(len(scales)):
            # Img = cv2.resize(img, (int(h*scales[k]), int(w*scales[k])), interpolation=cv2.INTER_CUBIC)
            Img = np.expand_dims(img, 2)
            patches = Im2Patch(Img, win=patch_size, stride=stride)
            print("file: %s scale %.1f # samples: %d" % (files[i], scales[k], patches.shape[3]))
            for n in range(patches.shape[3]):
                data_n = patches[:, :, :, n].copy()
                h5f.create_dataset(str(train_num), data=data_n)
                train_num += 1
    h5f.close()

    print('%s set, # samples %d\n' % (filename, train_num))


class Dataset(udata.Dataset):
    def __init__(self, train=True):
        super(Dataset, self).__init__()
        self.train = train
        self.transform = transforms.Compose([transforms.ToTensor(), ])  # range [0, 255] -> [0.0,1.0]
        if self.train:
            h5f = h5py.File('train.h5', 'r')
        else:
            h5f = h5py.File('val.h5', 'r')
        self.keys = list(h5f.keys())
        random.shuffle(self.keys)
        h5f.close()

    def __len__(self):
        return len(self.keys)

    def __getitem__(self, index):
        if self.train:
            h5f = h5py.File('train.h5', 'r')
        else:
            h5f = h5py.File('val.h5', 'r')
        key = self.keys[index]
        data = np.array(h5f[key])
        h5f.close()
        data = self.transform(data)
        img_data = {'data': data}
        return img_data








